#教与学算法
import random

import numpy as np
import math
import copy
import time

class TLBO():
    def __init__(self, n_dim=None, pop_size=20, max_iter=50, lb=-512, ub=512, obj_func=None):
        self.dim = n_dim
        self.pop = pop_size
        self.ub = ub
        self.lb = lb
        self.Max_iter = max_iter
        self.fun = obj_func
        self.X = np.zeros([pop_size, n_dim])
        for i in range(pop_size):
            for j in range(n_dim):
                self.X[i, j] = (ub[j] - lb[j]) * np.random.random() + lb[j]

    def BorderCheck(self):
        #边界检查
        pop = self.pop
        dim = self.dim
        lb = self.lb
        ub = self.ub
        for i in range(pop):
            for j in range(dim):
                if self.X[i,j] > ub[j]:
                    self.X[i, j] = ub[j]
                elif self.X[i,j] < lb[j]:
                    self.X[i, j] = lb[j]

    '''计算适应度函数'''

    def CaculateFitness(self):
        pop = self.X.shape[0]
        fitness = np.zeros([pop, 1])
        for i in range(pop):
            fitness[i] = self.fun(self.X[i, :])
        return fitness

    '''对适应度排序'''

    def SortFitness(self, Fit):
        fitness = np.sort(Fit, axis=0)
        index = np.argsort(Fit, axis=0)
        return fitness, index

    '''根据适应度对位置进行排序'''

    def SortPosition(self, index):
        Xnew = np.zeros(self.X.shape)
        for i in range(self.X.shape[0]):
            Xnew[i, :] = self.X[index[i], :]
        return Xnew

    def Tent_TLBO(self):
        pop = self.pop
        dim = self.dim
        ub = self.ub
        lb = self.lb
        fun = self.fun
        maxIter = self.Max_iter
        avg_fit = []  # 每次迭代的所有个体适应度平均值
        avg_ts = []  # 平均方差
        fitness = self.CaculateFitness()
        GbestScore = np.min(fitness)  #寻找最优适应度
        indexBest = np.argmin(fitness)  #最优适应度对应的索引
        GbestPosition = np.zeros([1,dim])
        GbestPosition[0,:] = copy.copy(self.X[indexBest,:]) #记录最优解
        Curve = np.zeros([maxIter,1])
        time_start = time.time()  # 记录迭代寻优开始时间
        for t in range(maxIter):
            for i in range(pop):
                #教阶段
                Xmean = np.mean(self.X)    #计算平均位置
                indexBest = np.argmin(fitness)  #寻找最优位置
                Xteacher = copy.copy(self.X[indexBest,:])  #教师的位置，也就是最优位置
                beta = random.randint(0,1)#教学因子
                Xnew = self.X[i,:] + np.random.random(dim) * (Xteacher - beta*Xmean)  #教阶段的位置更新

                #边界检查
                for j in range(dim):
                    if Xnew[j] > ub[j]:
                        Xnew[j] = ub[j]
                    if Xnew[j] < lb[j]:
                        Xnew[j] = lb[j]
                #计算新位置的适应度
                fitnessNew = fun(Xnew)
                #若新位置更优，则更新先前解
                if fitnessNew < fitness[i]:
                    self.X[i,:] = copy.copy(Xnew)
                    fitness[i] = copy.copy(fitnessNew)

                #学阶段
                p = random.randint(0,dim-1) # 随机选择一个索引
                while i == p:  #确保随机选择的索引不等于当前索引
                    p = random.randint(0,dim - 1)

                #学阶段的位置更新
                if fitness[i] < fitness[p]:
                    Xnew = self.X[i,:] + np.random.random(dim)* (self.X[i,:] - self.X[p,:])
                else:
                    Xnew = self.X[i,:] - np.random.random(dim)* (self.X[i,:] - self.X[p,:])

                # 边界检查
                for j in range(dim):
                    if Xnew[j] > ub[j]:
                        Xnew[j] = ub[j]
                    if Xnew[j] < lb[j]:
                        Xnew[j] = lb[j]
                #若新位置更优，则更新先前解
                fitnessNew = fun(Xnew)
                if fitnessNew < fitness[i]:
                    self.X[i,:] = copy.copy(Xnew)
                    fitness[i] = fitnessNew
            fitness = self.CaculateFitness()
            avg_fit.append(np.mean(fitness))
            avg_ts.append(np.var(fitness))
            indexBest = np.argmin(fitness)
            if fitness[indexBest] <= GbestScore:   #更新全局最优
                GbestScore = copy.copy(fitness[indexBest])
                GbestPosition[0,:] = copy.copy(self.X[indexBest,:])
            Curve[t] = GbestScore
        time_end = time.time()  # 记录迭代结束时间
        print(f'TLBO共花费 {time_end - time_start} 秒')
        return Curve, avg_fit, avg_ts


